import pdfplumber
import re


def extract_text_from_academic_pdf(pdf_path, txt_path):
    """
    从学术论文PDF中提取文本并优化格式
    :param pdf_path: PDF文件路径
    :param txt_path: 输出TXT文件路径
    """
    full_text = ""

    with pdfplumber.open(pdf_path) as pdf:
        for i, page in enumerate(pdf.pages):
            # 提取页面文本（优化参数处理学术论文格式）
            text = page.extract_text(
                x_tolerance=3,  # 水平间距容忍度
                y_tolerance=3,  # 垂直间距容忍度
                layout=True,  # 保持布局结构
                use_text_flow=True  # 保持文本流顺序
            )

            if text:
                # 应用文本格式优化
                text = optimize_academic_text(text, i)
                full_text += text + "\n\n"

    # 写入文本文件
    with open(txt_path, "w", encoding="utf-8") as f:
        f.write(full_text)


def optimize_academic_text(text, page_num):
    """
    优化学术论文文本格式
    :param text: 原始文本
    :param page_num: 页码
    :return: 优化后的文本
    """
    # 修复错误的换行（连接被断开的单词）
    text = re.sub(r"(\w)-\s+(\w)", r"\1\2", text)

    # 修复多栏文本中的错误换行
    if page_num > 0:  # 跳过标题页
        text = re.sub(r"(\w)\s*\n\s*(\w)", r"\1 \2", text)

    # 特殊格式处理
    text = re.sub(r"\$\{\}^{\text{(\w+)}}\$", r"^{\1}", text)  # 作者标记
    text = re.sub(r"(\d)\s+例", r"\1例", text)  # 数字连接
    text = re.sub(r"P\s*<\s*0\.0\s*5", r"P<0.05", text)  # 统计标记

    # 处理表格中的特殊符号
    text = re.sub(r"(\d)\s*±\s*(\d)", r"\1±\2", text)

    # 保留标题格式
    if page_num == 0:
        text = re.sub(r"·综合研究·", "\n·综合研究·\n", text)
        text = re.sub(r"摘要\s*目的", "摘要\n目的", text)

    # 处理参考文献格式
    if "参考文献" in text:
        text = re.sub(r"($$\d+$$)\s*([A-Z])", r"\1 \2", text)

    return text


# 使用示例
if __name__ == "__main__":
    # input_pdf = "床上瑜伽运动在人工髋关节置换病人运动康复中的应用_宋云雅 - 副本.pdf"
    # output_txt = "论文内容.txt"
    input_pdf = "髋关节置换术后定向肌群引导训练_华莉.pdf"
    output_txt = "髋关节置换术后定向肌群引导训练_华莉.txt"
    extract_text_from_academic_pdf(input_pdf, output_txt)
    print(f"文本提取完成，已保存至: {output_txt}")
